{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "08486c78",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "from plotly.subplots import make_subplots\n",
    "from scipy.optimize import least_squares\n",
    "from IPython.display import display \n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a76806f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 一、参数定义（保持不变）\n",
    "class StellarObject:\n",
    "    def __init__(self, name, l, b, distance, pm_l, pm_b, rv, is_pulsar=False):\n",
    "        self.name = name\n",
    "        self.l = np.deg2rad(l)\n",
    "        self.b = np.deg2rad(b)\n",
    "        self.distance = distance  # pc\n",
    "        self.pm_l = pm_l * 1e-3 / 3600 / 180 * np.pi  # rad/年\n",
    "        self.pm_b = pm_b * 1e-3 / 3600 / 180 * np.pi\n",
    "        self.rv = rv  # km/s\n",
    "        self.is_pulsar = is_pulsar\n",
    "\n",
    "    def get_position(self, t=0):\n",
    "        new_distance = self.distance + self.rv * t / 1022  # pc/年\n",
    "        x = new_distance * np.cos(self.b) * np.cos(self.l)\n",
    "        y = new_distance * np.cos(self.b) * np.sin(self.l)\n",
    "        z = new_distance * np.sin(self.b)\n",
    "        self.l += self.pm_l * t  # 更新角度\n",
    "        self.b += self.pm_b * t\n",
    "        return np.array([x, y, z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c5f40670",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 二、天体定义（保持不变）\n",
    "pulsars = [\n",
    "    StellarObject(\"PSR J0030+0451\", 12.3, 4.7, 130, 2.1, -0.5, 15, True),\n",
    "    StellarObject(\"PSR B1919+21\", 53.2, -7.1, 250, -1.8, 1.2, -22, True),\n",
    "    StellarObject(\"PSR J1744-1134\", 352.1, -15.6, 400, 3.5, 0.0, 30, True),\n",
    "]\n",
    "target_star = StellarObject(\"目标母星\", 28.5, 3.2, 200, 1.5, 0.8, 5)\n",
    "planet_orbit = {\n",
    "    \"a\": 2.0,\n",
    "    \"e\": 0.1,\n",
    "    \"i\": np.deg2rad(10),\n",
    "    \"Omega\": np.deg2rad(45),\n",
    "    \"omega\": np.deg2rad(30),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "396f5a1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 三、定位算法（保持不变）\n",
    "def planet_position(t):\n",
    "    M = 2 * np.pi * t\n",
    "    E = M + planet_orbit[\"e\"] * np.sin(M)\n",
    "    x = planet_orbit[\"a\"] * (np.cos(E) - planet_orbit[\"e\"])\n",
    "    y = planet_orbit[\"a\"] * np.sqrt(1 - planet_orbit[\"e\"] ** 2) * np.sin(E)\n",
    "    return np.array([x, y, 0])\n",
    "\n",
    "\n",
    "def calculate_parallax(star_pos, observer_pos):\n",
    "    baseline = observer_pos[1] - observer_pos[0]  # AU→pc\n",
    "    vector = star_pos - observer_pos[0]\n",
    "    return (\n",
    "        np.arctan2(np.linalg.norm(np.cross(baseline, vector)), np.dot(baseline, vector))\n",
    "        * 1e3\n",
    "        * 3600\n",
    "        * 180\n",
    "        / np.pi\n",
    "    )\n",
    "\n",
    "\n",
    "def residual(params, observations, pulsars_pos):\n",
    "    l, b, d = params\n",
    "    pos = d * np.array([np.cos(b) * np.cos(l), np.cos(b) * np.sin(l), np.sin(b)])\n",
    "    residuals = []\n",
    "    for i, (pl, pb) in enumerate(observations):\n",
    "        vec_target = pos - pulsars_pos[i]\n",
    "        vec_pulsar = pulsars_pos[i]\n",
    "        angle = np.arccos(\n",
    "            np.dot(vec_target, vec_pulsar)\n",
    "            / (np.linalg.norm(vec_target) * np.linalg.norm(vec_pulsar))\n",
    "        )\n",
    "        residuals.append(angle - np.sqrt(pl**2 + pb**2))\n",
    "    return residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "72c03799",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 四、可视化（使用FigureWidget和ipywidgets实现动态交互）\n",
    "def create_visualization():\n",
    "    # 使用FigureWidget替代普通Figure\n",
    "    fig = go.FigureWidget(make_subplots(\n",
    "        rows=1, cols=2,\n",
    "        specs=[[{'type': 'scene'}, {'type': 'scene'}]],\n",
    "        subplot_titles=(\"银河系坐标系（3D）\", \"天球投影（2D）\")\n",
    "    ))\n",
    "    \n",
    "    time_steps = list(range(0, 101, 10))\n",
    "    # 添加时间滑块\n",
    "    time_slider = widgets.IntSlider(\n",
    "        value=0,\n",
    "        min=0,\n",
    "        max=100,\n",
    "        step=10,\n",
    "        description='时间（年）:',\n",
    "        continuous_update=False  # 仅在释放滑块时更新\n",
    "    )\n",
    "    \n",
    "    # 初始化天体位置\n",
    "    for obj in pulsars + [target_star]:\n",
    "        fig.add_scatter3d(\n",
    "            x=[obj.get_position(0)[0]], y=[obj.get_position(0)[1]], z=[obj.get_position(0)[2]],\n",
    "            mode='markers', name=obj.name, \n",
    "            marker=dict(size=5, color='gray' if obj.is_pulsar else 'blue'),\n",
    "            row=1, col=1\n",
    "        )\n",
    "    \n",
    "    # 添加定位结果的空轨迹\n",
    "    fig.add_scatter3d(\n",
    "        x=[], y=[], z=[], mode='markers+text', name='定位结果',\n",
    "        marker=dict(size=8, color='red'), textposition='top center',\n",
    "        row=1, col=1\n",
    "    )\n",
    "    \n",
    "    # 定义按钮点击回调\n",
    "    def on_button_click(b):\n",
    "        # 获取当前时间\n",
    "        t = time_slider.value\n",
    "        # 计算脉冲星位置\n",
    "        pulsars_pos = np.array([p.get_position(t) for p in pulsars])\n",
    "        # 地球两个位置（相隔1AU）\n",
    "        earth_pos1 = np.array([1, 0, 0]) * 4.848e-6  # 1AU=4.848e-6 pc\n",
    "        earth_pos2 = np.array([-1, 0, 0]) * 4.848e-6\n",
    "        \n",
    "        observations = []\n",
    "        for i in range(3):\n",
    "            # 计算目标星与地球位置的关系\n",
    "            vec1 = target_star.get_position(t) - earth_pos1\n",
    "            vec2 = target_star.get_position(t) - earth_pos2\n",
    "            # 计算夹角（单位：毫角秒）\n",
    "            angle1 = np.arccos(np.dot(vec1, pulsars_pos[i]) / (np.linalg.norm(vec1) * np.linalg.norm(pulsars_pos[i])))\n",
    "            angle2 = np.arccos(np.dot(vec2, pulsars_pos[i]) / (np.linalg.norm(vec2) * np.linalg.norm(pulsars_pos[i])))\n",
    "            observations.append((np.degrees(angle1)*3600*1000, np.degrees(angle2)*3600*1000))\n",
    "        \n",
    "        # 最小二乘拟合\n",
    "        initial_guess = [target_star.l, target_star.b, target_star.distance]\n",
    "        result = least_squares(residual, initial_guess, args=(observations, pulsars_pos))\n",
    "        l_fit, b_fit, d_fit = result.x\n",
    "        \n",
    "        # 更新定位结果的位置\n",
    "        with fig.batch_update():\n",
    "            fig.data[-1].x = [d_fit * np.cos(b_fit) * np.cos(l_fit)]\n",
    "            fig.data[-1].y = [d_fit * np.cos(b_fit) * np.sin(l_fit)]\n",
    "            fig.data[-1].z = [d_fit * np.sin(b_fit)]\n",
    "            fig.data[-1].text = [f\"误差: {np.linalg.norm(result.fun):.2f} mas\"]\n",
    "    \n",
    "    # 创建按钮并绑定事件\n",
    "    button = widgets.Button(description=\"启动定位\")\n",
    "    button.on_click(on_button_click)\n",
    "    \n",
    "    # 组合控件布局\n",
    "    controls = widgets.VBox([button, time_slider])\n",
    "    dashboard = widgets.HBox([controls, fig])\n",
    "    \n",
    "    return dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "77bc5c66",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0600c54c4a1a463e853c16839493c745",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(VBox(children=(Button(description='启动定位', style=ButtonStyle()), IntSlider(value=0, continuous_u…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dashboard = create_visualization()\n",
    "display(dashboard)  # 使用IPython的display"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
